from ultralytics import YOLO
import cv2
import os
import json
from pathlib import Path
import numpy as np

# 创建输出目录
output_dir = Path("dataset")
images_dir = output_dir / "images"
labels_dir = output_dir / "labels"
visualization_dir = output_dir / "visualization"
images_dir.mkdir(parents=True, exist_ok=True)
labels_dir.mkdir(parents=True, exist_ok=True)
visualization_dir.mkdir(parents=True, exist_ok=True)

# 加载模型
model = YOLO("yolo11x-pose.pt")

# 打开视频文件
video_path = "match.mp4"
cap = cv2.VideoCapture(video_path)

frame_count = 0
while cap.isOpened():
    success, frame = cap.read()
    if not success:
        break
    
    # 每10帧处理一次，避免生成太多相似的图片
    if frame_count % 10 == 0:
        # 进行预测
        results = model(frame)
        
        # 获取关键点数据
        keypoints = results[0].keypoints.data
        if len(keypoints) > 0:  # 只保存检测到人的帧
            # 保存原始图片
            image_path = images_dir / f"frame_{frame_count:06d}.jpg"
            cv2.imwrite(str(image_path), frame)
            
            # 保存可视化标注图片
            vis_path = visualization_dir / f"frame_{frame_count:06d}.jpg"
            annotated_frame = results[0].plot()
            cv2.imwrite(str(vis_path), annotated_frame)
            
            # 准备标注数据
            annotations = []
            for person_keypoints in keypoints:
                # 转换为相对坐标
                h, w = frame.shape[:2]
                keypoints_normalized = []
                for kp in person_keypoints:
                    x, y = kp[0].item() / w, kp[1].item() / h
                    conf = kp[2].item() if len(kp) > 2 else 1.0
                    keypoints_normalized.extend([x, y, conf])
                
                # 创建标注数据
                annotation = {
                    "image_id": frame_count,
                    "category_id": 0,  # 0 代表人
                    "keypoints": keypoints_normalized,
                    "num_keypoints": len(keypoints_normalized) // 3
                }
                annotations.append(annotation)
            
            # 保存标注数据
            label_path = labels_dir / f"frame_{frame_count:06d}.json"
            with open(label_path, 'w') as f:
                json.dump({
                    "image_id": frame_count,
                    "image_path": str(image_path),
                    "visualization_path": str(vis_path),
                    "annotations": annotations
                }, f, indent=2)
            
            print(f"Saved frame {frame_count} with {len(annotations)} person(s)")
    
    frame_count += 1

cap.release()
print("Processing completed!")

# 创建数据集配置文件
dataset_config = {
    "train": str(images_dir),
    "val": str(images_dir),  # 这里可以设置验证集路径
    "nc": 1,  # 类别数量
    "names": ["person"],  # 类别名称
    "visualization_dir": str(visualization_dir)
}

with open(output_dir / "dataset.yaml", 'w') as f:
    import yaml
    yaml.dump(dataset_config, f, default_flow_style=False)

print("Dataset configuration file created!")